#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：data_reader 
@File ：DatabaseReader.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2024/3/19 16:04 
"""

import pandas as pd
from sqlalchemy import create_engine, inspect


class DatabaseReader:
	def __init__(self):
		self.db_type = None
		self.db_params = {}
		self.table_name = None
		self.field_mapping = {}
		self.field_types = {}

	def set_database_type(self, db_type):
		"""设置数据库类型"""
		self.db_type = db_type

	def set_database_params(self, db_params):
		"""设置或更新数据库连接参数"""
		self.db_params = db_params

	def read_data_with_mapping(self, table_name, field_mapping={}, field_types={}):
		"""读取并筛选表数据"""
		self.table_name = table_name
		self.field_mapping = field_mapping
		self.field_types = field_types

		# 根据数据库类型创建连接引擎
		if self.db_type == 'mysql':
			engine_url = f'mysql+pymysql://{self.db_params["user"]}:{self.db_params["password"]}@{self.db_params["host"]}/{self.db_params["database"]}'
		elif self.db_type == 'sqlserver':
			engine_url = f'mssql+pyodbc://{self.db_params["user"]}:{self.db_params["password"]}@{self.db_params["host"]}/{self.db_params["database"]}?driver=ODBC+Driver+17+for+SQL+Server'
		elif self.db_type == 'oracle':
			engine_url = f'oracle+cx_oracle://{self.db_params["user"]}:{self.db_params["password"]}@{self.db_params["host"]}:{self.db_params["port"]}/{self.db_params["service_name"]}'
		else:
			raise ValueError("不支持的数据库类型，请选择mysql、sqlserver或oracle。")

		engine = create_engine(engine_url)

		# 筛选字段名及转换字段名称
		column_names = []
		for db_column in inspect(engine).get_columns(table_name):
			column_name = db_column['name']
			mapped_name = field_mapping.get(column_name, column_name)
			column_names.append(mapped_name)

		# 根据field_types字典添加类型转换
		type_conversions = {mapped_name: getattr(pd, field_types.get(mapped_name, 'str')) for mapped_name in
							column_names}

		query = f"SELECT {' ,'.join(column_names)} FROM {table_name}"
		df = pd.read_sql_query(query, engine)

		# 应用类型转换
		df = df.astype(type_conversions)

		return df


if __name__ == '__main__':
	print("DatabaseReader")
	# 使用示例
	db_reader = DatabaseReader()
	db_reader.set_database_type('sqlserver')
	db_reader.set_database_params({
		"user": "sa",
		"password": "123456",
		"host": "127.0.0.1",
		"database": "ColdRollingFormal"
	})

	field_map = {"original_field1": "new_field1", "original_field2": "new_field2"}
	field_types = {"new_field1": "int", "new_field2": "float"}

	# data = db_reader.read_data_with_mapping("surf_index", field_mapping=field_map, field_types=field_types)
	data = db_reader.read_data_with_mapping("surf_index")

	print(data)
